import pydicom
import pydicom.uid
import os
import nibabel as nib
import scipy.io as sio
import numpy as np
import cv2
from libtiff import TIFF


def nii_mask_dcm(dcm_folder, nii_path, save_path):
    # 读取nii文件
    nii = nib.load(nii_path)
    nii_arr = nii.get_fdata().astype(np.uint16)  # 16位色图

    for root, dirs, files in os.walk(dcm_folder):
        idx = 0
        for f in files:
            dcm_path = os.path.join(root, f)
            # 读取dicom文件，原文件是1通道灰度图
            dcm = pydicom.read_file(dcm_path, force=True)
            dcm.file_meta.TransferSyntaxUID = pydicom.uid.ImplicitVRLittleEndian
            dcm_arr = dcm.pixel_array   # 像素值转化成ndarray

            # 每层nii
            nii_slice = nii_arr[:, :, idx].transpose(1, 0)  # 调整nii和dcm方向一致，idx层数和dcm文件个数一致

            _, nii_mask = cv2.threshold(nii_slice, 1, 65535, cv2.THRESH_BINARY)  # 二值化
            img_masked = np.bitwise_and(dcm_arr, nii_mask)  # mask掩码操作，按位与，ndarray

            # 排除掉mask后无效(全黑)的图片，像素值=0
            if np.any(img_masked):
                np.save(os.path.join(save_path, '{}.npy'.format(idx)), img_masked)   # 保存numpy数组
                sio.savemat(os.path.join(save_path, '{}.mat'.format(idx)), {'masked': img_masked})   # matlab文件
                tiff = TIFF.open(os.path.join(save_path, '{}.tif'.format(idx)), mode='w')    # 无损压缩的tif图片
                tiff.write_image(img_masked, compression=None)
                tiff.close()
            idx += 1


if __name__ == '__main__':

    dcm_folder_9614225_1 = 'To_SJC_gastric_data/9614225/dicom/20090917/372025/1.2.840.113619.2.55.3.12588800.5666.1253143047.577'
    dcm_folder_9614225_2 = 'To_SJC_gastric_data/9614225/dicom/20091109/383682/1.2.840.113619.2.55.3.12588800.5808.1257720980.307'
    dcm_folder_9625145_1 = 'To_SJC_gastric_data/9625145/dicom/20090825/365972/1.2.840.113619.2.55.3.12588800.5720.1251155538.300'
    dcm_folder_9625145_2 = 'To_SJC_gastric_data/9625145/dicom/20091010/375532/1.3.12.2.1107.5.1.1.20862.20091010085028893.3'
    dcm_folder_9625145_3 = 'To_SJC_gastric_data/9625145/dicom/20091010/375532/1.3.12.2.1107.5.1.1.20862.20091010085210653.3'
    dcm_folders = [dcm_folder_9614225_1, dcm_folder_9614225_2, dcm_folder_9625145_1, dcm_folder_9625145_2]  # dcm文件路径

    nii_path_9614255_1 = 'To_SJC_gastric_data/9614225/label/9614225-1-label.nii'
    nii_path_9614255_2 = 'To_SJC_gastric_data/9614225/label/9614225-2-label.nii'
    nii_path_9625145_1 = 'To_SJC_gastric_data/9625145/label/9625145-1-label.nii'
    nii_path_9625145_2 = 'To_SJC_gastric_data/9625145/label/9625145-2-label.nii'
    nii_paths = [nii_path_9614255_1, nii_path_9614255_2, nii_path_9625145_1, nii_path_9625145_2]

    save_path = './dcm_masked_nii/'
    save_next_path = ['9614225_1/', '9614225_2/', '9625145_1/', '9625145_2/']  # 保存具体分类文件夹

    i = 0
    for folder in dcm_folders:
        save_paths = os.path.join(save_path, save_next_path[i])
        if not os.path.exists(save_paths):
            os.makedirs(save_paths)

        nii_mask_dcm(folder, nii_paths[i], save_paths)
        i += 1

    # 145病人另一个文件夹也是23，单独处理
    # save_paths = os.path.join(save_path, '9625145_3/')
    # nii_mask_dcm(dcm_folder_9625145_3, nii_paths[3], save_paths)








